/*
 * To change this license header, choose License Headers in Project Properties.
 * To change this template file, choose Tools | Templates
 * and open the template in the editor.
 */
package com.compomics.dataprovider.model.fasta;

import com.compomics.dataprovider.control.logic.fasta.FastaSourceConverter;
import com.compomics.util.experiment.identification.FastaIndex;
import com.compomics.util.experiment.identification.SequenceFactory;
import com.compomics.util.protein.Header;
import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.TreeMap;
import org.apache.log4j.Logger;

/**
 *
 * @author Kenneth Verheggen
 */
public class FastaStream {

    private final String tag = "<COMPOMICS_DATA_PROCESSING>";
    private ByteArrayInputStream fastaAsStream;
    private static final Logger LOGGER = Logger.getLogger(FastaStream.class);
    private int entries;
    private int decoys;
    private boolean saved = false;

    /**
     *
     * @param inputStream the starting point for this stream
     */
    public FastaStream(ByteArrayInputStream inputStream) {
        this.fastaAsStream = inputStream;
    }

    /**
     *
     * @return the fasta as a stream
     */
    public ByteArrayInputStream toStream() {
        return fastaAsStream;
    }

    //INPUT METHODS--------------------------------------------------------------------------------------------------
    /**
     *
     * @return a the byteArrayOutputStream starting from the given input
     * @throws IOException
     */
    public ByteArrayOutputStream getOutputStream() throws IOException {
        ByteArrayOutputStream out;
        out = new ByteArrayOutputStream();
        int i;
        while ((i = toStream().read()) != -1) {
            out.write(i);
            out.flush();
        }
        return out;
    }

    //CONCATENATION METHODS--------------------------------------------------------------------------------------------------   
    /**
     *
     * @param fastaSource
     * @return a new Fastastream with the given source concatenated
     * @throws IOException
     */
    public FastaStream concatenate(Object fastaSource) throws IOException {
        LOGGER.debug("Concatenating streams");
        ByteArrayInputStream inputStream = FastaSourceConverter.convertToStream(fastaSource);
        ByteArrayOutputStream outputStream = getOutputStream();
        try (BufferedWriter psout = new BufferedWriter(new OutputStreamWriter(outputStream)); BufferedReader in = new BufferedReader(new InputStreamReader(inputStream))) {
            String inputLine;
            psout.newLine();
            int lines = 0;
            while ((inputLine = in.readLine()) != null) {
                inputLine = inputLine.replaceAll("\\r|\\n", "");
                if (inputLine.startsWith(">")) {
                    if (lines > 0) {
                        psout.write(System.lineSeparator());
                    }
                    psout.write(inputLine + System.lineSeparator());
                } else {
                    psout.write(inputLine);
                }
                lines++;
            }
        }
        return commit(outputStream);
    }

    //FINALISATION METHODS--------------------------------------------------------------------------------------------------
    private FastaStream commit(ByteArrayOutputStream result) {
        LOGGER.debug("Updating FastaStream");
        fastaAsStream = FastaSourceConverter.redirectOutputStream(result);
        return this;
    }

    /**
     * @return FastaStream with appended decoys
     * @throws IOException
     */
    public FastaStream addFastDecoys() throws IOException {
        LOGGER.debug("Adding decoys without using sequence factory");
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        try (BufferedWriter psout = new BufferedWriter(new OutputStreamWriter(outputStream)); BufferedReader in = new BufferedReader(new InputStreamReader(toStream()))) {
            String inputLine;
            StringBuilder currentdecoyHeader = new StringBuilder();
            int lines = 0;
            boolean shortHeader;
            while ((inputLine = in.readLine()) != null) {
                inputLine = inputLine.replaceAll("\\r|\\n", "");
                String temp = inputLine;
                if (temp.startsWith(">")) {
                    if (lines > 0) {
                        psout.write(System.lineSeparator());
                    }
                    psout.write(temp + System.lineSeparator());
                    //change the decoy header
                    String[] split = inputLine.split("\\|");
                    currentdecoyHeader = new StringBuilder(System.lineSeparator()).append(split[0]).append("|").append(split[1]).append("_REVERSED").append("|");
                    shortHeader = split.length < 2;
                    if (!shortHeader) {
                        try {
                            currentdecoyHeader.append(split[2]);
                        } catch (ArrayIndexOutOfBoundsException e) {
                            shortHeader = true;
                        }
                    }
                    if (shortHeader) {
                        currentdecoyHeader.append(split[1]);
                    }
                    currentdecoyHeader.append("-REVERSED").append(System.lineSeparator());
                } else {
                    psout.write(inputLine);
                    psout.write(currentdecoyHeader.toString());
                    psout.write(new StringBuilder(inputLine).reverse().toString());
                }
                lines++;
            }
        }
        return commit(outputStream);
    }

    //FINALISATION METHODS
    //--------------------------------------------------------------------------------------------------
    /**
     * @return the updated UniprotFastaStream
     * @throws IOException
     */
    public FastaStream addDecoys() throws IOException {
        LOGGER.debug("Adding decoys using sequence factory");
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        try (BufferedWriter psout = new BufferedWriter(new OutputStreamWriter(outputStream)); BufferedReader in = new BufferedReader(new InputStreamReader(toStream()))) {
            String inputLine;
            String decoyHeader;
            StringBuilder sequence;
            while ((inputLine = in.readLine()) != null) {
                if (inputLine.startsWith(">")) {
                    sequence = new StringBuilder(in.readLine());
                    decoyHeader = getDecoyHeader(inputLine);
                    psout.write(inputLine + System.lineSeparator());
                    psout.write(sequence.toString() + System.lineSeparator());
                    psout.write(decoyHeader + System.lineSeparator());
                    psout.write(sequence.reverse().toString() + System.lineSeparator());
                    psout.flush();
                }
            }
            return commit(outputStream);
        }
    }

    private String getDecoyHeader(String header) {
        Header decoyHeader = Header.parseFromFASTA(header);
        String defaultDecoyAccession = SequenceFactory.getDefaultDecoyAccession(decoyHeader.getAccession());
        String defaultDecoyDescription = SequenceFactory.getDefaultDecoyDescription(decoyHeader.getDescription());
        decoyHeader.setAccession(defaultDecoyAccession);
        decoyHeader.setDescription(defaultDecoyDescription);
        return decoyHeader.toString();
    }

    /**
     *
     * @return the updated UniprotFastaStream
     * @throws IOException
     */
    public FastaStream appendCRAP() throws IOException, NullPointerException {
        LOGGER.debug("Adding common contaminants");
        File crapDb = new File(FastaStream.class.getResource("/crap.fasta").getFile());
        return concatenate(FastaSourceConverter.convertToStream(crapDb));
    }

    //POST PROCESSING METHODS--------------------------------------------------------------------------------------------------
    /**
     * WARNING : this method will only keep the first encountered accession in
     * case of duplicates
     *
     * @return the updated UniprotFastaStream
     * @throws IOException
     */
    public FastaStream removeDuplicateAccessions() throws IOException {
        LOGGER.debug("Removing duplicate accessions");
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        try (BufferedWriter psout = new BufferedWriter(new OutputStreamWriter(outputStream)); BufferedReader in = new BufferedReader(new InputStreamReader(toStream()))) {
            String inputLine;
            HashSet<String> encounteredAccessions = new HashSet<>();
            String accession;
            String sequence;
            while ((inputLine = in.readLine()) != null) {
                accession = inputLine.split("\\|")[1];
                sequence = in.readLine();
                if (!encounteredAccessions.contains(accession)) {
                    encounteredAccessions.add(accession);
                    psout.write(inputLine + System.lineSeparator());
                    psout.write(sequence + System.lineSeparator());
                    psout.flush();
                }
            }
        }
        return commit(outputStream);
    }

    /**
     * WARNING : this method will use the first encountered accession as the
     * main entry, equivalents will be noted in the description part of the
     * header
     *
     * @return the updated UniprotFastaStream
     * @throws IOException
     */
    public FastaStream removeDuplicateSequences() throws IOException {
        LOGGER.debug("Removing duplicate sequences");
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        try (BufferedWriter psout = new BufferedWriter(new OutputStreamWriter(outputStream)); BufferedReader in = new BufferedReader(new InputStreamReader(toStream()))) {
            String header;
            String sequence;
            HashSet<String> encounteredAccessions = new HashSet<>();
            HashMap<String, String> encounteredSequencesWithHeaders = new HashMap<>();
            String headersForSequence;
            String accession;
            //read sequences and headers...
            while ((header = in.readLine()) != null) {
                String[] headerParts = header.split("\\|");
                if (headerParts.length > 1) {
                    accession = headerParts[1];
                    //read the sequences
                    sequence = in.readLine();
                    if (!encounteredAccessions.contains(accession) && sequence != null) {
                        encounteredAccessions.add(accession);
                        if (encounteredSequencesWithHeaders.containsKey(sequence)) {
                            headersForSequence = encounteredSequencesWithHeaders.get(sequence);
                            if (!headersForSequence.contains(tag)) {
                                String firstAccession = headersForSequence.split("\\|")[1];
                                int endIndex = headersForSequence.lastIndexOf("|");
                                headersForSequence = headersForSequence.substring(0, endIndex) + "|" + tag + firstAccession;
                                if (!headersForSequence.endsWith("\\|")) {
                                    headersForSequence += " -AND- ";
                                }
                            }
                            headersForSequence += accession;
                        } else {
                            headersForSequence = header;
                        }
                        encounteredSequencesWithHeaders.put(sequence, headersForSequence);
                    }
                }
            }
            for (String aSequence : encounteredSequencesWithHeaders.keySet()) {
                //make new combined header...
                if (aSequence.length() >= 5) {
                    String headers = encounteredSequencesWithHeaders.get(aSequence);
                    psout.write(headers.replace(tag, "") + System.lineSeparator());
                    psout.flush();
                    psout.write(aSequence + System.lineSeparator());
                    psout.flush();
                }
            }
        }
        return commit(outputStream);
    }

    /**
     * WARNING : this method will keep all duplicates, but rename them case of
     * duplicates
     *
     * @return the updated UniprotFastaStream
     * @throws IOException
     */
    @Deprecated
    public FastaStream renameDuplicateAccessions() throws IOException {
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        try (BufferedWriter psout = new BufferedWriter(new OutputStreamWriter(outputStream)); BufferedReader in = new BufferedReader(new InputStreamReader(toStream()))) {
            String inputLine;
            HashSet<String> encounteredAccessions = new HashSet<>();
            String accession;
            String temp;
            while ((inputLine = in.readLine()) != null) {
                inputLine = inputLine.replaceAll("\\r|\\n", "");
                if (inputLine.startsWith(">")) {
                    temp = accession = inputLine.split("\\|")[1];
                    int i = 2;
                    while (encounteredAccessions.contains(temp)) {
                        temp = temp.substring(0, accession.length()) + "-" + i;
                        i++;
                    }
                    encounteredAccessions.add(temp);
                    if (encounteredAccessions.size() > 0) {
                        psout.write(System.lineSeparator());
                    }
                    psout.write(inputLine + System.lineSeparator());
                } else {
                    psout.write(inputLine);
                }
            }
        }
        return commit(outputStream);
    }

    /**
     * WARNING : this method will use the first encountered accession as the
     * main entry, equivalents will be noted in the description part of the
     * header
     *
     * @return the updated UniprotFastaStream
     * @throws IOException
     */
    @Deprecated
    public FastaStream sortOnAccession() throws IOException {
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        try (BufferedWriter psout = new BufferedWriter(new OutputStreamWriter(outputStream)); BufferedReader in = new BufferedReader(new InputStreamReader(toStream()))) {
            String header;
            String sequence;
            TreeMap<String, String> encounteredSequencesWithHeaders = new TreeMap<>();
            //read sequences and headers...
            while ((header = in.readLine()) != null) {
                sequence = in.readLine();
                encounteredSequencesWithHeaders.put(header, sequence);
            }
            for (Map.Entry<String, String> aHeader : encounteredSequencesWithHeaders.entrySet()) {
                psout.write(aHeader.getKey() + System.lineSeparator());
                psout.write(aHeader.getValue() + System.lineSeparator());
            }
        }
        return commit(outputStream);
    }

    /**
     * WARNING : this method will use the first encountered accession as the
     * main entry, equivalents will be noted in the description part of the
     * header
     *
     * @return the updated UniprotFastaStream
     * @throws IOException
     */
    @Deprecated
    public FastaStream sortOnSequenceLength() throws IOException {
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        try (BufferedWriter psout = new BufferedWriter(new OutputStreamWriter(outputStream)); BufferedReader in = new BufferedReader(new InputStreamReader(toStream()))) {
            String header;
            String sequence;
            TreeMap<Integer, String[]> encounteredSequencesWithHeaders = new TreeMap<>();
            //read sequences and headers...
            while ((header = in.readLine()) != null) {
                sequence = in.readLine();
                encounteredSequencesWithHeaders.put(sequence.length(), new String[]{header, sequence});
            }
            String[] entry;
            for (Map.Entry<Integer, String[]> aHeader : encounteredSequencesWithHeaders.entrySet()) {
                entry = aHeader.getValue();
                psout.write(entry[0] + System.lineSeparator());
                psout.write(entry[1] + System.lineSeparator());
            }
        }
        return commit(outputStream);
    }

    /**
     *
     * @param outputFile the file where the fasta should be saved to
     * @param fastaFileName the theoretical file name for the fasta in the
     * stream
     * @param fastaName the theoretical name for the fasta in the stream
     * @param append
     * @return the compomics FastaIndex object for this stream
     * @throws IOException
     */
    public FastaIndex indexAndSave(File outputFile, String fastaFileName, String fastaName, boolean append) throws IOException {
        LOGGER.debug("Indexing and Saving FASTA stream to file");
        try (BufferedWriter out = new BufferedWriter(new FileWriter(outputFile, append)); BufferedReader in = new BufferedReader(new InputStreamReader(toStream()))) {
            String inputLine;
            HashMap<String, Long> indices = new HashMap<>();
            HashSet<String> decoys = new HashSet<>();
            long bytePosition = 0L;
            String accession;
            while ((inputLine = in.readLine()) != null) {
                bytePosition = bytePosition + inputLine.getBytes().length;
                if (inputLine.startsWith(">")) {
                    accession = inputLine.split("\\|")[1];
                    indices.put(accession, bytePosition);
                    if (accession.endsWith("REVERSED")) {
                        decoys.add(accession);
                    }
                }
                out.write(inputLine + System.lineSeparator());
            }
            return new FastaIndex(indices, decoys, fastaFileName, fastaName,
                    true,
                    true,
                    indices.size() - decoys.size(),
                    System.currentTimeMillis(),
                    Header.DatabaseType.Unknown,
                    "REVERSED",
                    FastaIndex.getDefaultVersion(System.currentTimeMillis())
            );
        }
    }

    /**
     *
     * @param fastaFileName the theoretical file name for the fasta in the
     * stream
     * @param fastaName the theoretical name for the fasta in the stream
     * @return the compomics FastaIndex object for this stream
     * @throws IOException
     */
    public FastaIndex index(String fastaFileName, String fastaName) throws IOException {
        LOGGER.debug("Indexing FASTA stream");
        try (BufferedReader in = new BufferedReader(new InputStreamReader(toStream()))) {
            String inputLine;
            HashMap<String, Long> indices = new HashMap<>();
            HashSet<String> decoys = new HashSet<>();
            long bytePosition = 0L;
            String accession;
            while ((inputLine = in.readLine()) != null) {
                bytePosition = bytePosition + inputLine.getBytes().length;
                if (inputLine.startsWith(">")) {
                    accession = inputLine.split("\\|")[1];
                    indices.put(accession, bytePosition);
                    if (accession.endsWith("REVERSED")) {
                        decoys.add(accession);
                    }
                }
            }
            return new FastaIndex(indices, decoys, fastaFileName, fastaName,
                    true,
                    true,
                    indices.size() - decoys.size(),
                    System.currentTimeMillis(),
                    Header.DatabaseType.Unknown,
                    "REVERSED",
                    FastaIndex.getDefaultVersion(System.currentTimeMillis())
            );
        }
    }

    /**
     * @param outputFile the file where the fasta should be saved to
     * @param append
     * @throws IOException
     */
    public void save(File outputFile, boolean append) throws IOException {
        try (FileWriter out = new FileWriter(outputFile, append); BufferedReader in = new BufferedReader(new InputStreamReader(toStream()))) {
            String inputLine;
            while ((inputLine = in.readLine()) != null) {
                out.append(inputLine).append(System.lineSeparator());
                out.flush();
                if (inputLine.startsWith(">")) {
                    entries++;
                    if (inputLine.contains("_REVERSED")) {
                        decoys++;
                    }
                }
            }
        }
        saved = true;
    }

    /**
     *
     * @return the amount of entries AFTER saving
     */
    public int getEntries() throws IllegalStateException{
        if (isSaved()) {
            return entries;
        } else {
            throw new IllegalStateException("The stream has not been saved yet !");
        }
    }

    /**
     *
     * @return the amount of entries AFTER saving
     */
    public int getDecoys() throws IllegalStateException{
        if (isSaved()) {
            return decoys;
        } else {
            throw new IllegalStateException("The stream has not been saved yet !");
        }
    }

    /**
     *
     * @return boolean indicating whether the stream has been saved already
     */
    public boolean isSaved() {
        return saved;
    }

}
